# -*- coding: utf-8 -*-

# -------------------------------
# Author    : linlingwei
# Email     : 951035650@qq.com
# File Name : spIkhandle.py
# Date      : 2019/12/11 
# IDE       : PyCharm
# Version   : 1.0.1
# -------------------------------

import maya.cmds as mc
import maya.mel as mel

import configs.constants.attribute_constant as ac
import configs.constants.nodeType_constant as nc
import configs.constants.suffix_constant as sc
import configs.constants.body_constant as bc
import configs.constants.global_constant as gc

import scripts.commonCmds as con
import scripts.curveCmds as crc

import rig.ikCore as ikCore
reload(ikCore)

class SpIkhandle(ikCore.IkCore):
    def __init__(self,lname="test"):
        super(SpIkhandle,self).__init__("ikSplineSolver",lname=lname)
    
    
    def create(self,sjot,ejot,crv=None):
        self._init()
        if not crv:
            self.ikhandle(sjot, ejot,createCurve=True)
        else:
            self.ikhandle(sjot, ejot,crv)
            # mc.ikHandle(self._ikhandle,e=True,c=crv)
        self.initJoint()
        return

    

    def stretch(self):
        jotList = self.jotlist
        star_poci = None
        ikCurve_shape = con.findRelatedShape(self.ikCurve)
        for i, jnt in enumerate(jotList[:-1], 1):
            end_poci = mc.createNode(nc.tPointOnCurveInfo, name=con.lname(jotList[i], sc.sPoci))
            mc.connectAttr(ikCurve_shape + ".ws", end_poci + ".ic")
            par = con.getObjInCrvParameter(self.ikCurve, jotList[i])
            mc.setAttr(end_poci + '.pr', par)
        
            if not star_poci:
                star_poci = mc.createNode(nc.tPointOnCurveInfo, name=con.lname(jnt, sc.sPoci))
                mc.connectAttr(ikCurve_shape + ".ws", star_poci + ".ic")
                par = con.getObjInCrvParameter(self.ikCurve, jnt)
                mc.setAttr(star_poci + '.pr', par)
        
            dis_dbn = mc.createNode(nc.tDistanceBetween, name=con.lname(jnt, sc.sDbn))
            mc.connectAttr(star_poci + ".p", dis_dbn + ".p1")
            mc.connectAttr(end_poci + ".p", dis_dbn + ".p2")
        
            local_md = mc.createNode(nc.tMultiplyDivide, name=con.lname(jnt, ac.local_scale, sc.sMd))
            mc.connectAttr(dis_dbn + ".d", local_md + ".i1x")
            mc.connectAttr(con.lAttr(self.localBridge, ac.global_scale), local_md + ".i2x")
            mc.setAttr(local_md + ".op", 2)
            
            mc.connectAttr(local_md+".ox",con.lAttr(jnt,ac.stretch))
            # refresh *star_poci
            star_poci = end_poci
        
        return
    
    def initJoint(self, axial='x'):
        if 'init' in mc.listAttr(self.localBridge):
            return
        mc.addAttr(self.localBridge, ln='init', at='bool', dv=True, k=False)
        con.lockHideAttr(self.localBridge, ['init'], lock=True, hide=True)
        if ac.global_scale not in mc.listAttr(self.localBridge):
            mc.addAttr(self.localBridge, ln=ac.global_scale, at='float', dv=1, k=True)
        if ac.local_scale not in mc.listAttr(self.localBridge):
            mc.addAttr(self.localBridge, ln=ac.local_scale, at='float', dv=1, min=0.01, k=True)
        attr = ac.stretch_switch
        if attr not in mc.listAttr(self.localBridge):
            mc.addAttr(self.localBridge, ln=attr, at='float', dv=0, min=0, max=1, k=True)
        attr = ac.keep_volume_switch
        if attr not in mc.listAttr(self.localBridge):
            mc.addAttr(self.localBridge, ln=attr, at='float', dv=1, min=0, max=1, k=True)
        attr = ac.keep_volume_ratio
        if attr not in mc.listAttr(self.localBridge):
            mc.addAttr(self.localBridge, ln=attr, at='float', dv=0, k=True)
            
        jotList = self.jotlist
        for i,jnt in enumerate(jotList[:-1],1):
            jntLen = mc.getAttr(jotList[i] + '.t%s' % axial)
            attr = ac.local_length
            if attr not in mc.listAttr(jnt):
                mc.addAttr(jnt, ln=attr, at='float', dv=jntLen)
                mc.setAttr('{}.{}'.format(jnt, attr), jntLen, l=True)
    
            attr = ac.offset
            if attr not in mc.listAttr(jnt):
                mc.addAttr(jnt, ln=attr, at='float', dv=0)
                mc.setAttr('{}.{}'.format(jnt, attr), 0, k=True, )
            attr = ac.stretch
            if attr not in mc.listAttr(jnt):
                mc.addAttr(jnt, ln=attr, at='float', dv=0)
                mc.setAttr('{}.{}'.format(jnt, attr), 0, k=True, )
            jntLenPma = mc.createNode('plusMinusAverage',
                                      name=con.lname(jnt, ac.local_length, sc.sPma))

            mc.setAttr(jntLenPma + '.operation', 1)

            # mc.connectAttr(jnt + '.{}'.format(ac.local_length), jntLenPma + '.input1D[0]')
            mc.connectAttr(jnt + '.' + ac.offset, jntLenPma + '.input1D[0]')
            # mc.connectAttr(jnt + '.' + ac.stretch, jntLenPma + '.input1D[2]')
            
            length_bta = mc.createNode(nc.tBlendTwoAttr,name=con.lname(jnt,ac.local_length,sc.sBta))
            mc.connectAttr(con.lAttr(jnt,ac.local_length), length_bta + '.input[0]')
            mc.connectAttr(con.lAttr(jnt,ac.stretch), length_bta + '.input[1]')
            mc.connectAttr(con.lAttr(self.localBridge,ac.stretch_switch),length_bta+".ab")
            mc.connectAttr(length_bta+".o",jntLenPma + '.input1D[1]')
            
            mc.connectAttr(jntLenPma + '.output1D', jotList[i] + '.t%s' % axial)

            attr = ac.keep_volume
            if attr not in mc.listAttr(jnt):
                mc.addAttr(jnt, ln=attr, at='float', dv=1)
                mc.setAttr('{}.{}'.format(jnt, attr), 1, k=True, )
            
            # ratio_adl = mc.createNode(nc.tAddDoubleLinear,name=con.lname(jnt,ac.keep_volume_ratio,sc.sMdl))
            # mc.connectAttr(con.lAttr(self.localBridge,ac.keep_volume_ratio),ratio_adl+".i1")
            # mc.connectAttr(con.lAttr(jnt,ac.keep_volume),ratio_adl+".i2")
            
            kv_bta = mc.createNode(nc.tBlendTwoAttr,name=con.lname(jnt,ac.keep_volume_switch,sc.sBta))
            mc.connectAttr(con.lAttr(self.localBridge,ac.keep_volume_switch),kv_bta+".ab")
            mc.connectAttr(con.lAttr(jnt,ac.keep_volume),kv_bta+".input[1]")
            mc.setAttr(kv_bta+".input[0]",1)
            mc.connectAttr(kv_bta+".o",jnt+".sy")
            mc.connectAttr(kv_bta + ".o", jnt + ".sz")
            

        return

    def advRotate(self,  hide=True):
        ikhandle = self._ikhandle
        jotlist = self.jotlist
    
        starJnt = jotlist[0]
        effect = jotlist[-1]
        mc.setAttr(effect+".jointOrient",0,0,0)
    
        starLoc = con.objToLocator(starJnt, con.lname(self.lname, "AdvRot_Star"), parent=False, hide=hide)
        endLoc = con.objToLocator(effect, con.lname(self.lname, "AdvRot_End"), parent=False, hide=hide)
    
        mc.setAttr(ikhandle + ".dTwistControlEnable", 1)
        mc.setAttr(ikhandle + ".dWorldUpType", 4)
        mc.setAttr(ikhandle + ".dForwardAxis", 0)
        mc.setAttr(ikhandle + ".dWorldUpAxis", 0)
        mc.setAttr(ikhandle + ".dWorldUpVector", 0, 1, 0)
        mc.setAttr(ikhandle + ".dWorldUpVectorEnd", 0, 1, 0)
    
        mc.connectAttr(starLoc + ".worldMatrix[0]", ikhandle + ".dWorldUpMatrix")
        mc.connectAttr(endLoc + ".worldMatrix[0]", ikhandle + ".dWorldUpMatrixEnd")
    
        return starLoc, endLoc
    
    def keepVolumn(self):
        jotlist = self.jotlist
        for i,jot in enumerate(jotlist[:-1]):
            kv_md = mc.createNode(nc.tMultiplyDivide,name = con.lname(jot,ac.keep_volume,sc.sMd))
            mc.connectAttr(con.lAttr(jot,ac.local_length),kv_md+".i1x")
            mc.connectAttr(con.lAttr(jot,ac.stretch),kv_md+".i2x")
            mc.setAttr(kv_md+".op",2)
            # mc.connectAttr(kv_md+".ox",con.lAttr(jot,ac.keep_volume))
            
            kv_pma = mc.createNode(nc.tPlusMinusAverage,name=con.lname(jot,ac.keep_volume,sc.sPma))
            mc.setAttr(kv_pma+".input1D[1]",1)
            mc.connectAttr(kv_md+".ox",kv_pma+".input1D[0]")
            mc.setAttr(kv_pma+".op",2)
            
            ratio_mdl = mc.createNode(nc.tMultDoubleLinear,name=con.lname(jot,ac.keep_volume_ratio,sc.sMdl))
            mc.connectAttr(kv_pma+".output1D",ratio_mdl+".i1")
            mc.connectAttr(con.lAttr(self.localBridge,ac.keep_volume_ratio),ratio_mdl+".i2")
            
            ratio_adl = mc.createNode(nc.tAddDoubleLinear,name=con.lname(jot,ac.keep_volume_ratio,sc.sAdl))
            mc.connectAttr(kv_md+".ox",ratio_adl+".i1")
            mc.connectAttr(ratio_mdl+".o",ratio_adl+".i2")
            
            mc.connectAttr(ratio_adl+".o",con.lAttr(jot,ac.keep_volume))
            
        return

if __name__ == '__main__':
    pass